
from socketserver import BaseRequestHandler, TCPServer, ThreadingTCPServer
import random
import flag
import os
import string
from hashlib import sha256
import signal

def s2b(s):
    res = b''
    for i in s:
        res += bytes([ord(i)])
    return res

def b2s(b):
    s = ''
    for i in b:
        s += chr(i)
    return s

class Hash():

    def __init__(self):
        self.h = 144066263297769815596495629667062367629
        self.g = (1 << 88) + 315
        self.mask = 0xffffffffffffffffffffffffffffffff
    
    def encode(self,message):
        message = [ord(i) for i in message]
        h = self.h
        for i in message:
            h = ((h+i)*self.g)&self.mask
        return h

class Task(BaseRequestHandler):


    def dosend(self, msg):
        try:
            self.request.sendall(msg)
        except:
            pass

    def recvall(self, sz):
        try:
            r = sz
            res = ""
            while r > 0:
                res += b2s(self.request.recv(r))
                if res.endswith("\n"):
                    r = 0
                else:
                    r = sz - len(res)
            res = res.strip()
        except:
            res = ""
        return res.strip("\n")

    def proof_of_work(self):
        random.seed(os.urandom(8))
        proof = ''.join([random.choice(string.ascii_letters+string.digits) for _ in range(20)])
        digest = sha256(proof.encode()).hexdigest()
        self.request.send(s2b(("sha256(XXXX+%s) == %s\n" % (proof[4:],digest))))
        self.request.send(b'Give me XXXX:')
        x = self.request.recv(10)
        x = b2s(x)
        x = x.strip()
        if len(x) != 4 or sha256((x+proof[4:]).encode()).hexdigest() != digest: 
            return False
        return True

    def handle(self):
        
        signal.alarm(300)
        if not self.proof_of_work():
            return
        signal.alarm(300)

        self.dosend(b"welcome to eZLogin System!\n")
        admin = "admin_is_badmonkey"
        hash = Hash()
        sig = str(hash.encode(admin))
        self.dosend("admin's signature is: ")
        self.dosend(s2b(sig+'\n'))
        try:
            while True:
                self.dosend(b"Options\n")
                self.dosend(b"1.Login\n")
                self.dosend(b"2.Register\n")
            
                choice = int(self.recvall(1024))
                if choice != 1 and choice != 2:
                    self.dosend(b"no such choice\n")
                    break
                if choice == 1:
                    self.dosend(b"give me your username(hex)\n")
                    username = self.recvall(1024)
                    username = b2s(bytes.fromhex(username))
                    usersig = str(hash.encode(username))
                    print(username)
                    if username == admin :
                        self.dosend(b"sorry,you are not badmonkey\n")
                        break
                    if usersig == sig:
                        self.dosend(b"welcome,badmonkey\n")
                        self.dosend(flag)
                        break
                    else:
                        self.dosend(s2b("welcome,{}\n".format(username)))
                else:
                    self.dosend(b"give me your username(hex)\n")
                    username = self.recvall(1024)
                    username = b2s(bytes.fromhex(username))

                    usersig = str(hash.encode(username))
                    self.dosend(b"Here is your signature\n")
                    self.dosend(s2b(usersig+'\n'))
        except e:
            print(e)
            self.dosend(b"wrong input\n")
            return False

if __name__ == '__main__':
    HOST, PORT = "183.129.189.62", 57202
    server = ThreadingTCPServer((HOST, PORT), Task)
    server.allow_reuse_address = True
    server.serve_forever()


